import numpy as np


# 混淆矩阵计算
def confusion_matrix(real, predict, n=2):
    '''
    :param real: 真实标签(batchsize, H, W)
    :param predict:  预测标签(batchsize, H, W)
    :return:  混淆矩阵
    '''
    real = real.to('cpu')
    predict = predict.to('cpu')
    real = real.long()
    predict = predict.long()
    batch_size = real.shape[0]
    CM = np.zeros(shape=(batch_size,n,n), dtype=np.int64)
    real_flatten = real.flatten(start_dim = 1)
    predict_flatten = predict.flatten(start_dim = 1)
    for i in range(batch_size):
        cm = n * real_flatten[i,:] + predict_flatten[i,:]
        cm = np.bincount(cm, minlength=n ** 2).reshape(n, n)
        CM[i,:, :] = cm
    return CM

# 像素精度
def pixelAccuracy(CM):
    '''
    :param CM: 混淆矩阵
    :return: 像素精度
    '''
    acc = np.mean(CM.diagonal(offset=0, axis1=1, axis2=2).sum(axis=1) / CM.sum(axis=1).sum(axis=1)).item()
    return acc

# 类别像素精度
def clsPixelAccuracy(CM):
    clsacc = CM.diagonal(0, 1, 2) / CM.sum(axis=2)
    clsacc  = np.nan_to_num(clsacc, nan = 0)
    return np.mean(clsacc, axis=0)

# MIOU
def MIOU(CM):
    '''
    :param CM: 混淆矩阵
    :return: 平均交并比
    '''
    iou = CM.diagonal(0,1,2)/(CM.sum(axis=2)+CM.sum(axis=1)-CM.diagonal(0,1,2))
    return np.mean(iou, axis=1).mean()

